import os
from typing import Annotated



from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId, tool
from langchain_openai import ChatOpenAI
from typing_extensions import TypedDict
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.types import Command, interrupt

# 定义一个状态字典
class State(TypedDict):
    messages: Annotated[list, add_messages]
    name: str
    birthday: str




# 定义一个工具函数
@tool
def human_assistance(
        name: str,
        birthday: str,
        tool_call_id: Annotated[str, InjectedToolCallId]
) -> str:

    """Request assistance from a human."""
    human_response = interrupt(
        {
            "question": "Is this correct?",
            "name": name,
            "birthday": birthday,
        },
    )

    if human_response.get("correct", "").lower().startswith("y"):
        verified_name = name
        verified_birthday = birthday
        response = "Correct"
    else:
        verified_name = human_response.get("name", name)
        verified_birthday = human_response.get("birthday", birthday)
        response = f"Made a correction: {human_response}"

    state_update = {
        "name": verified_name,
        "birthday": verified_birthday,
        "messages": [ToolMessage(response, tool_call_id=tool_call_id)],
    }

    return Command(update=state_update)





tool = TavilySearchResults(max_results=2)

tools = [tool, human_assistance]

llm = ChatOpenAI(model_name="gpt-4", base_url=os.environ["OPENAI_BASE_URL"])

llm_with_tools = llm.bind_tools(tools)





def chatbot(state: State):

    print("chatbot message:",state["messages"])
    message = llm_with_tools.invoke(state["messages"])

    assert(len(message.tool_calls) <= 1)

    return {"messages": [message]}





graph_builder = StateGraph(State)

graph_builder.add_node("chatbot", chatbot)



tool_node = ToolNode(tools=tools)

graph_builder.add_node("tools", tool_node)



graph_builder.add_conditional_edges(

    "chatbot",

    tools_condition,

)

graph_builder.add_edge("tools", "chatbot")

graph_builder.add_edge(START, "chatbot")



memory = MemorySaver()

graph = graph_builder.compile(checkpointer=memory)

graph.invoke({"messages":{
    "role": "user",
    "content": "my name is xiaoming and my birthday is 1990-01-01"
}},config={"configurable":{"thread_id":42}})

# 将生成的图片保存到文件中
graph_png = graph.get_graph().draw_mermaid_png()
with open("code14-1-chatbot.png", "wb") as f:
    f.write(graph_png)